/*
 * HkfModel.cpp
 *
 *  Created on: 1 Aug 2011
 *      Author: Allan
 */

#include "HkfModel.h"

// C++ includes
#include <cmath>
#include <map>
using namespace std;

// GeoReact includes
#include "Geochemistry/ElectrolyteSolution.h"
#include "Numerics/LinearInterpolator.h"
#include "Utils/SpeciesUtils.h"

const double ElectrostaticParameterEta = 1.66027E+05;

const double 
DebyeHuckelParameterA(double T, double P);

const double 
DebyeHuckelParameterB(double T, double P);

const double 
HkfParameterNaCl(double T, double P);

const double 
HkfParameterNapClm(double T, double P);

const double 
TotalMolalitySolutes(const ElectrolyteSolution& sol, const VectorXd& n);

const double HkfIonicActivityCoefficient(const ElectrolyteSolution& sol, double T, double P, const VectorXd& n, unsigned ion, const vector<double>& effRadii)
{
	// The electrical charges of the ionic species
	const VectorXd zi = sol.GetIonicCharges();
	
	// The effective ionic strength of the solution
	//const double I = sol.IonicStrengthEffective(n);
	const double I = sol.StoichiometricIonicStrength(n); // It seems that the stoichiometric ionic strength yields much more precise activity coefficients than the effective one
	
	// The square root of the ionic strength
	const double sqrtI = sqrt(I);
	
	// The total molality of the solutes
	const double mT = TotalMolalitySolutes(sol, n);
	
	// The parameters for the HKF model
	const double A       = DebyeHuckelParameterA(T, P);
	const double B       = DebyeHuckelParameterB(T, P);
	const double bNaCl   = HkfParameterNaCl(T, P);
	const double bNapClm = HkfParameterNapClm(T, P);
	
	// The electrical charge of the current ion
	const double z = zi[ion];

	// The square of the electrical charge of the ion
	const double z2 = z * z;

	// The effective electrostatic radius of the ion
	const double effRadius = effRadii[ion];

	// The absolute Born coefficient of the ion
	const double omegaAbs = ElectrostaticParameterEta * z2 / effRadius;

	// The Debye-Huckel ion size parameter of the ion as computed by Reed, 1982 and also in TOUGHREACT
	const double a = (z < 0) ? 2.0 * (effRadius + 1.91 * abs(z))/(abs(z) + 1.0) : 
							   2.0 * (effRadius + 1.81 * abs(z))/(abs(z) + 1.0);

	// The log10 of the activity coefficient of the ion in molality scale
	const double log_gamma = -(A * z2 * sqrtI)/(1.0 + a * B * sqrtI) - log10(1.0 + 0.0180153 * mT) + (omegaAbs * bNaCl + bNapClm - 0.19 * (abs(z) - 1.0)) * I;
	
	// Store the activity coefficient of the current ion
	return std::pow(10.0, log_gamma);
}

const double HkfWaterActivityCoefficient(const ElectrolyteSolution& sol, double T, double P, const VectorXd& n, const vector<double>& effRadii)
{
	// The electrical charges of the ionic species
	const VectorXd zi = sol.GetIonicCharges();
	
	// The number of ionic species
	const unsigned numIons = zi.size();
	
	// The effective ionic strength of the solution
	//const double I = sol.IonicStrengthEffective(n);
	const double I = sol.StoichiometricIonicStrength(n); // It seems that the stoichiometric ionic strength yields much more precise activity coefficients than the effective one
	
	// The square root of the ionic strength
	const double sqrtI = sqrt(I);
	
	// The total molality of the solutes
	const double mT = TotalMolalitySolutes(sol, n);
	
	// The parameters for the HKF model
	const double A       = DebyeHuckelParameterA(T, P);
	const double B       = DebyeHuckelParameterB(T, P);
	const double bNaCl   = HkfParameterNaCl(T, P);
	const double bNapClm = HkfParameterNapClm(T, P);

	// The Born coefficient of the ion H[+]
	const double omegaH = 0.5387E+05;
	
	// The stoichiometric molality of the ions
	const VectorXd smions = sol.StoichiometricIonicMolalities(n);
	
	// The osmotic coefficient of the aqueous phase
	double osmoticCoeff = 0.0;
	
	// Loop over all ions in the solution to compute the osmotic coefficient
	for(unsigned ion = 0; ion < numIons; ++ion)
	{
		// The electrical charge of the current ion species
		const double z = zi[ion];
		
		// The square of the electrical charge
		const double z2 = z * z;
		
		// The effective electrostatic radius of the current ion
		const double effRadius = effRadii[ion];
		
		// The Born coefficient of the current ion
		const double omega = ElectrostaticParameterEta * z2 / effRadius - z * omegaH;
		
		// The Debye-Huckel ion size parameter of the ion as computed by Reed, 1982 and also in TOUGHREACT
		const double a = (z < 0) ? 2.0 * (effRadius + 1.91 * abs(z))/(abs(z) + 1.0) : 
				                   2.0 * (effRadius + 1.81 * abs(z))/(abs(z) + 1.0);
		
		// The Lambda parameter of the current ion
		const double Lambda = 1.0 + a * B * sqrtI;
		
		// The sigma parameter of the current ion
		const double sigma = 3.0/pow(a * B * sqrtI, 3) * (Lambda - 1.0/Lambda - 2.0 * log(Lambda));
		
		// calculate the psi contribution of the current ion
		double psi = A * z2 * sqrtI * sigma/3.0 - log10(1.0 + 0.0180153 * mT)/(0.0180153 * mT) - 0.5 * (omega * bNaCl + bNapClm - 0.19 * (abs(z) - 1.0)) * I;
		
		osmoticCoeff += smions[ion] * psi;
	}
	
	// Correct the osmotic coefficient
	osmoticCoeff *= -2.303/mT;
	
	// Compute the activity of water
	const double activityWater = exp(-osmoticCoeff * mT/55.508);
	
	// Return the activity coefficient of water divided by 55.508, since we define gw = aw/55.508, where gw and aw are the activity coefficient and activity of water respectively
	return activityWater / 55.508;
}

std::map<string,double> HkfEffectiveRadii = 
{
	{"H[+]",    3.08},	{"Li[+]",   1.64},	{"Na[+]",   1.91},	{"K[+]",    2.27},
	{"Rb[+]",   2.41},	{"Cs[+]",   2.61},	{"NH4[+]",  2.31},	{"Ag[+]",   2.20},
	{"Au[+]",   2.31},	{"Cu[+]",   1.90},	{"Mg[2+]",  2.54},	{"Sr[2+]",  3.00},
	{"Ca[2+]",  2.87},	{"Ba[2+]",  3.22},	{"Pb[2+]",  3.08},	{"Zn[2+]",  2.62},
	{"Cu[2+]",  2.60},	{"Cd[2+]",  2.85},	{"Hg[2+]",  2.98},	{"Fe[2+]",  2.62},
	{"Mn[2+]",  2.68},	{"Fe[3+]",  3.46},	{"Al[3+]",  3.33},	{"Au[3+]",  3.72},
	{"La[3+]",  3.96},	{"Gd[3+]",  3.79},	{"In[3+]",  3.63},	{"Ca[3+]",  3.44},
	{"F[-]",    1.33},	{"Cl[-]",   1.81},	{"Br[-]",   1.96},	{"I[-]",    2.20},
	{"OH[-]",   1.40},	{"HS[-]",   1.84},	{"NO3[-]",  2.81},	{"HCO3[-]", 2.10},
	{"HSO4[-]", 2.37},	{"ClO4[-]", 3.59},	{"ReO4[-]", 4.23},	{"SO4[2-]", 3.15},
	{"CO3[2-]", 2.81}
};

LinearInterpolator DebyeHuckelA({     0,     25,     60,    100,    150,    200,    250,    300}, 
		                        {0.4939, 0.5114, 0.5465, 0.5995, 0.6855, 0.7994, 0.9593, 1.2180});

LinearInterpolator DebyeHuckelB({     0,     25,     60,    100,    150,    200,    250,    300}, 
		                        {0.3253, 0.3288, 0.3346, 0.6842, 0.3421, 0.3639, 0.3766, 0.3925});

LinearInterpolator HkfNaCl({  25,   50,   75,  100,  125,  150,   175,   200,   225,   250,   275,   300},
		                    {2.47, 2.15, 1.79, 1.39, 0.93, 0.41, -0.18, -0.85, -1.64, -2.57, -3.71, -5.21});

LinearInterpolator HkfNapClm({   25,    50,    75,  100,  125,  150,  175,  200,  225,   250,   275,   300},
		                     {-9.77, -5.59, -2.43, 0.23, 2.44, 4.51, 6.38, 8.11, 9.73, 11.29, 12.71, 14.15});

const double DebyeHuckelParameterA(double T, double P)
{
	return DebyeHuckelA(T);
}

const double DebyeHuckelParameterB(double T, double P)
{
	return DebyeHuckelB(T);
}

const double HkfParameterNaCl(double T, double P)
{
	const double R = 1.9858775;
	
	const double factor = 1.0E-3 * 1.0/(2.302585 * R * (T + 273.15));

	return factor * HkfNaCl(T);
}

const double HkfParameterNapClm(double T, double P)
{
	return 1.0E-2 * HkfNapClm(T);
}

const double EffectiveElectrostaticRadius(const string& ion)
{
	// Find the effective electrostatic radius of the given ion
	auto iterator = HkfEffectiveRadii.find(ion);

	// Check if the ion has been found in the list of ions
	if(iterator != HkfEffectiveRadii.end())
		return iterator->second;
	else
	{
		// The charge of the ion
		const int charge = (int)ElectricalCharge(ion);

		// The approach below was taken from Table H.1 of:
		// Xu, T., Sonnenthal, E., Spycher, N., & Pruess, K. (2004). TOUGHREACT-User's Guide: A Simulation Program for Non-isothermal Multiphase 
		// Reactive geochemical Transport in Variable Saturated Geologic Media. Computers & Geosciences (Vol. 32).
		switch(charge)
		{
		case -1: return 1.81; // Assuming the Cl[-] value
		case -2: return 3.00; // Assuming the rounded average of CO3[2-] and SO4[2-] values
		case -3: return 4.20; // Estimated from straight line fit with charge
		case +1: return 2.31; // Assuming the NH4[+] value
		case +2: return 2.80; // Assuming the rounded average for +2 species in HKF Table 3
		case +3: return 3.60; // Assuming the rounded average for +3 species in HKF Table 3
		case +4: return 4.50; // Estimated using HKF Equation (unknown eq. number) and average crystallographic radii of +4 species in CRC Handbook
		default:
			if(charge > +3)
				return std::abs(charge) * 4.5/4.0; // Linear Extrapolation (charge � 4.5/3.0) 
			else // (charge < -3)
				return std::abs(charge) * 4.2/3.0; // Linear Extrapolation (charge � 4.2/4.0)
		}
	}
}

const double TotalMolalitySolutes(const ElectrolyteSolution& sol, const VectorXd& n)
{
	return 55.508 * n.sum() / n[sol.GetWaterIndex()] - 55.508;
}
